"""
网络请求监控脚本
用于记录所有 HTTP/HTTPS 请求和响应信息
"""
from mitmproxy import ctx
import json
from datetime import datetime
import os
import sys

# 获取脚本所在目录
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
# 日志文件路径
LOG_FILE = os.path.join(SCRIPT_DIR, 'network_log.txt')


def load(loader):
    print(f"日志文件将保存在: {LOG_FILE}")
    print("正在监听所有网络请求...")

    # 创建或清空日志文件
    with open(LOG_FILE, 'w', encoding='utf-8') as f:
        f.write(f"网络监控日志 - 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write("监听范围: 所有网络请求\n")
        f.write("=" * 50 + "\n")
        f.flush()


def save_to_log(content):
    """保存内容到日志文件并立即刷新"""
    try:
        with open(LOG_FILE, 'a', encoding='utf-8', buffering=1) as f:  # buffering=1 启用行缓冲
            f.write(content)
            f.flush()  # 立即刷新缓冲区
            os.fsync(f.fileno())  # 确保写入磁盘
        print(f"成功写入日志: {content[:100]}...")  # 只打印前100个字符
    except Exception as e:
        print(f"写入日志失败: {str(e)}")


def request(flow):
    try:
        request_info = {
            "时间": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "URL": flow.request.pretty_url,
            "方法": flow.request.method,
            "主机": flow.request.pretty_host,
            "端口": flow.request.port,
            "协议": flow.request.scheme,
            "请求头": dict(flow.request.headers),
            "请求体": flow.request.content.decode('utf-8', errors='ignore') if flow.request.content else None
        }

        log_content = f"\n{'=' * 50}\n请求信息:\n{json.dumps(request_info, ensure_ascii=False, indent=2)}\n"
        save_to_log(log_content)
        print(f"记录了请求: {flow.request.pretty_url}")

    except Exception as e:
        error_msg = f"处理请求时出错: {str(e)}\n"
        print(error_msg)
        save_to_log(f"\n错误信息:\n{error_msg}")


def response(flow):
    try:
        response_info = {
            "时间": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "URL": flow.request.pretty_url,
            "主机": flow.request.pretty_host,
            "状态码": flow.response.status_code,
            "响应头": dict(flow.response.headers),
            "内容类型": flow.response.headers.get("content-type", ""),
            "响应体": flow.response.content.decode('utf-8', errors='ignore') if flow.response.content else None
        }

        log_content = f"\n{'=' * 50}\n响应信息:\n{json.dumps(response_info, ensure_ascii=False, indent=2)}\n"
        save_to_log(log_content)
        print(f"记录了响应: {flow.request.pretty_url} (状态码: {flow.response.status_code})")

    except Exception as e:
        error_msg = f"处理响应时出错: {str(e)}\n"
        print(error_msg)
        save_to_log(f"\n错误信息:\n{error_msg}")


def done():
    """脚本结束时的处理"""
    try:
        save_to_log(f"\n{'=' * 50}\n监控结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n{'=' * 50}\n")
    except Exception as e:
        print(f"记录结束时间失败: {str(e)}")


# 启动时打印提示信息
print(f"""
{'=' * 50}
网络监控脚本已启动
监听范围: 所有网络请求
脚本位置: {__file__}
日志文件: {LOG_FILE}

注意: 监听 HTTPS 请求需要安装 mitmproxy 证书
安装步骤:
1. 确保 mitmproxy 正在运行
2. 访问 http://mitm.it
3. 下载并安装对应系统的证书
{'=' * 50}
""")
